#例题2.15
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
# 求解范围
xa, xb = 0.0, 1.0  # 空间范围，xa<x<xb
ya, yb = 0.0, 1.0  # 空间范围，ya<y<yb
# 初始化
h = 0.01  # 空间步长, dx = dy = 0.01
w = 0.5  # 松弛因子
nodes = round((xb-xa)/h)  # $\mathrm{x}$轴 空间网格数
# 边值条件
u = np.zeros((nodes+1, nodes+1))
for i in range(nodes+1):
    u[i, 0] = 1.0 + np.sin(0.5*(i-50)/np.pi)
    u[i, -1] = -1.0 + 0.5*np.sin((i-50)/np.pi)
    u[0, i] = -1.0 + 0.5*np.sin((i-50)/np.pi)
    u[-1, i] = 1.0 + np.sin(0.5*(50-i)/np.pi)
# 迭代松弛法求解
for iter in range(100):
    for i in range(1, nodes):
        for j in range(1, nodes):
            u[i, j] = w/4 * (u[i-1, j] + u[i+1, j] + u[i, j-1] + u[i, j+1]) + (1-w) * u[i, j]
# 绘图
x = np.linspace(0, 1, nodes+1)
y = np.linspace(0, 1, nodes+1)
xx, yy = np.meshgrid(x, y)
fig = plt.figure(figsize=(8,6))
ax = fig.add_subplot(111, projection='3d')
surf = ax.plot_surface(xx, yy, u, cmap=plt.get_cmap('rainbow'))
fig.colorbar(surf, shrink=0.5)
ax.set_xlim3d(0, 1.0)
ax.set_ylim3d(0, 1.0)
ax.set_zlim3d(-2, 2.5)
ax.set_title("2D elliptic partial differential equation")
ax.set_xlabel("X")
ax.set_ylabel("Y")
plt.show()
